# 加载环境变量
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv())

import os
import json
from pathlib import Path
from pprint import pprint
import sqlite3

import sys
sys.path.append('..')
from Utils.PrintUtils import color_print, OBSERVATION_COLOR
from Utils.PromptTemplateBuilder import PromptTemplateBuilder

from langchain.prompts import ChatPromptTemplate
from langchain.prompts.chat import SystemMessagePromptTemplate, HumanMessagePromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI
from langchain.tools import StructuredTool


# db = ":memory:"
db = "./data/LM_APP_001.db"

def load_data(json_string_path="./data", 
              json_string_file="hotels.json", 
              database_string_path="./prompts/tools", 
              database_schema_string_file="database_schema_string.txt", 
              database_statement_string_file="database_statement_string.txt"):
    # 加载数据
    # loader = DirectoryLoader('./data', glob="**/*.json", 
    #                          loader_cls=JSONLoader, loader_kwargs={"jq_schema": '.', "text_content": False})
    # documents = loader.load()
    # print(f'加载文档数量：{len(documents)}')

    hotels = json.loads(Path(os.path.join(json_string_path, json_string_file)).read_text())

    # pprint(hotels[:3])
    # print("************************")

    # 创建数据库连接
    conn = sqlite3.connect(db)
    cur = conn.cursor()

    # 删除上次运行产生的hotels表
    cur.execute("drop table if exists hotels;")

    # 创建hotels表
    cur.execute(Path(os.path.join(database_string_path, database_schema_string_file)).read_text())

    # 插入hotels字典列表
    cur.executemany(
        Path(os.path.join(database_string_path, database_statement_string_file)).read_text(), 
        [(h["hotel_id"], h["name"], h["type"], h["address"], h["subway"], 
          h["phone"], h["facilities"], h["price"], h["rating"]) for h in hotels]
    )

    # 提交事务
    conn.commit()

    # try:
    #     cur.execute("select * from hotels;") # 所有查询结果都在cur对象里
    #     data2 = cur.fetchmany(3)
    #     pprint(data2)
    #     print("************************")
    #     conn.commit()
    # except Exception as e:
    #     print(e)

    cur.close()
    conn.close()

class DatabaseQuerist():
    def __init__(self,
                 prompts_path="./prompts/tools",
                 prompt_file_1="database_querist_1.json",
                 prompt_file_2="database_querist_2.json",
                 verbose=False,
                 web_ui=False
    ):
        self.prompt_1 = PromptTemplateBuilder(prompts_path, prompt_file_1).build()
        self.prompt_2 = PromptTemplateBuilder(prompts_path, prompt_file_2).build()
        self.verbose = verbose
        self.web_ui = web_ui

    def query(self, query: str):
        """查询数据库，回答用户问题。"""
        llm = ChatOpenAI(
            # model="gpt-4-1106-preview",
            model="deepseek-chat",
            temperature=0,
            seed=42
        )

        chain = {"query": RunnablePassthrough()} | self.prompt_1 | llm | StrOutputParser()

        # sql_query = chain.invoke(query)
        if self.verbose:
            s = f"SQL query: "
            color_print(s, OBSERVATION_COLOR, end="\n")
            if self.web_ui:
                yield s + "\n"
        sql_query = ""
        for s in chain.stream(query):
            sql_query += s
            if self.verbose:
                color_print(s, OBSERVATION_COLOR, end="")
                if self.web_ui:
                    yield s
        if self.verbose:
            color_print("", OBSERVATION_COLOR, end="\n")
            if self.web_ui:
                yield "\n"

        # 创建数据库连接
        conn = sqlite3.connect(db)
        cur = conn.cursor()

        data = None
        sql_query_result = None
        result = None

        # 查询hotels表
        try:
            cur.execute(sql_query) # 所有查询结果都在cur对象里

            data = cur.fetchall()

            num = len(data)
            if num == 0:
                data = None
                result = "未找到相关酒店信息。"
                if self.verbose:
                    s = f"\n----\n结果：\n{result}"
                    color_print(s, OBSERVATION_COLOR)
                    if self.web_ui:
                        yield s + "\n"
            elif num <= 10:
                # pprint(data)
                sql_query_result = f"SQL查询结果有{num}条记录，依次为：{data}。"
            else:
                # pprint(data[:10])
                sql_query_result = f"SQL查询结果有{num}条记录，前10条记录依次为：{data[:10]}。"

            conn.commit()
        except Exception as e:
            if self.verbose:
                color_print(e, OBSERVATION_COLOR)
                if self.web_ui:
                    yield e + "\n"
            raise e

        if data is not None:
            llm = ChatOpenAI(
                # model="gpt-4-1106-preview",
                model="deepseek-chat",
                temperature=0,
                seed=42
            )

            chain = {"query": RunnablePassthrough(), "sql_query": RunnablePassthrough(), "sql_query_result": RunnablePassthrough()} \
                  | self.prompt_2 | llm | StrOutputParser()

            # result = chain.invoke((query, sql_query, sql_query_result))
            if self.verbose:
                s = f"\n----\n结果："
                color_print(s, OBSERVATION_COLOR, end="\n")
                if self.web_ui:
                    yield s + "\n"
            result = ""
            for s in chain.stream((query, sql_query, sql_query_result)):
                result += s
                if self.verbose:
                    color_print(s, OBSERVATION_COLOR, end="")
                    if self.web_ui:
                        yield s

        cur.close()
        conn.close()

        return result

    def as_tool(self):
        return StructuredTool.from_function(
            func=self.query,
            name="QueryDatabase",
            description="使用自然语言查询数据库，回答用户问题。",
        )


if __name__ == "__main__":
    os.chdir("..")  # 将当前工作目录更改为上一级目录

    load_data() # 加载数据

    # query = "500-800元价格范围内的酒店"
    # query = "200元以下的酒店"
    # query = "600元以上的酒店"
    # query = "酒店名称为鹏润酒店"
    # query = "在广安门附近的酒店"
    # query = "在北京西站地铁站附近的酒店"
    # query = "评分不低于4.5的酒店"
    # query = "有健身房和WiFi的酒店"
    # query = "有健身房或游泳池的酒店"
    query = "需同时满足以下几点：1. 一家高档型酒店。2. 最低价格在600-1000元之间的酒店。3. 一间有游泳池的酒店。"
    print(f"query: \n{query}")

    result = DatabaseQuerist(verbose=True).query(query)
    print(f"result: \n{result}")
